import json
import os.path
import re
from typing import List


class BaseSplitter:
    def __init__(self, book_path) -> None:
        self.book_path = book_path
        with open(book_path, "r", encoding="utf-8") as f:
            self.text = f.read()

        self.chapter_pattern = None
        self.section_pattern = None
        self.sub_sec_pattern = None
        self.level_4_pattern = None
        self.define_regex()

        self.output = []
        self.LEN_LIMIT_HIGH = 2500
        self.LEN_LIMIT_LOW = 1000

    def define_regex(self):
        """
        this function should be implemented in child class
        it's used to define regex patterns for chapter, section and sub-section or deeper levels
        """
        raise NotImplementedError

    def _preprocess(self, content):
        """
        remove img element in markdown: ![.*?](.*?)
        remove table element, its in html format
        replace multiple empty lines with single one
        """
        content = re.sub(r"!\[.*?\]\(.*?\)", "", content)
        content = re.sub(r"<table>.*?</table>", "", content, flags=re.DOTALL)
        content = re.sub(r"\n{2,}", "\n", content)
        return content

    def _split(self, content, pattern):
        """
        a common function to split content by pattern and retrieve all titles
        """

        parts = re.split(pattern, content)

        titles = re.findall(pattern, content)
        # titles = [title.strip("# ") for title in titles]
        tmp = []
        for title in titles:
            tmp.append(title.strip("# "))
            if "双分量模型" in title:
                print(title)

        assert len(parts) == len(titles) + 1, "something seems wrong with spliting content by pattern"

        # create a list to store all seperated parts
        results = [{"content": part.strip(), "title": title} for part, title in zip(parts[1:], titles)]
        # insert at 0 the first preface part
        if parts[0].strip() != "":
            results.insert(0, {"content": parts[0].strip(), "title": ""})

        return results

    def _add_element(self, content):
        if len(content) > 1.5 * self.LEN_LIMIT_HIGH:
            self._devide_large_part(content.split("\n"))
        else:
            self.output.append(
                {
                    "id": len(self.output) + 1,
                    "chapter": self.chapter,
                    "section": self.section,
                    "sub_section": self.sub_section,
                    "content": content,
                }
            )

    def _devide_large_part(self, parts: List[str]):
        """
        split a large part into smaller parts to avoid exceeding the limit

        Args:
         - parts: List[str]
        """

        small_chunks = ""
        for i, p in enumerate(parts):
            small_chunks += p + "\n"
            # the residual part should not be too short
            if i == len(parts) - 1:
                if len(small_chunks) > self.LEN_LIMIT_LOW:
                    self._add_element(small_chunks)
                return
            if (
                len(small_chunks) <= self.LEN_LIMIT_HIGH
                and len(small_chunks + parts[i + 1]) >= self.LEN_LIMIT_HIGH
            ):
                self._add_element(small_chunks)
                small_chunks = ""

    def _try_split_by_level4(self, content):
        if not self.level_4_pattern:
            return False

        results = self._split(content, self.level_4_pattern)
        if len(results) > 1:
            chunks = [result["title"] + "\n" + result["content"] for result in results]
            self._devide_large_part(chunks)
            return True
        else:
            return False

    def _try_split_by_title(self, content):
        title_pattern = re.compile(r"^##.+?$", re.M)
        results = self._split(content, title_pattern)
        if len(results) > 2:
            chunks = [result["title"] + "\n" + result["content"] for result in results]
            # print([result["title"] for result in results])
            self._devide_large_part(chunks)
            return True
        else:
            return False

    def _try_split_by_paragraph(self, content: str):
        paragraphs = content.split("\n")
        self._devide_large_part(paragraphs)

    def add_new_content(self, content):
        """
        add some segemented text to output json
        if the text is too long, split it into multile parts
        by joining paragraphs until the length of content is greater than limit
        """

        if len(content) < self.LEN_LIMIT_HIGH:
            self._add_element(content)
        else:
            # if self._try_split_by_level4(content):
            #     return
            if self._try_split_by_title(content):
                return
            else:
                self._try_split_by_paragraph(content)

    def determine_level(self, line):
        """
        determine the level of the title line by counting the numbers in the title
        """
        title_nums = re.search(r"\d.+\d", line)
        if title_nums:
            nums = title_nums.group()
            if "." in nums:
                return len(nums.split("."))
            elif "-" in nums:
                return len(nums.split("-"))
            else:
                return len(re.findall(r"\d+", nums))
        else:
            return 0

    def process_book(self, save_dir):
        """
        iterate through each line and check if it is a chapter, section or sub-section
        """
        content = ""
        self.chapter, self.section, self.sub_section = "", "", ""
        text = self._preprocess(self.text)
        # text = self.text
        for line in text.split("\n"):
            if line.strip() == "":
                continue
            if (
                (chapter := self.chapter_pattern.search(line))
                or (section := self.section_pattern.search(line))
                # or (sub_section := self.sub_sec_pattern.match(line))
            ):
                if len(content) > self.LEN_LIMIT_LOW:
                    self.add_new_content(content)
                    content = line + "\n"

                if chapter:
                    self.chapter = chapter.group().strip("# ")
                    self.section, self.sub_section = "", ""
                elif section:
                    level = self.determine_level(line)
                    if level == 2:
                        self.section = section.group().strip("# ")
                        # self.section = re.sub(r"[^\d\.\-]", "", section)
                        self.sub_section = ""
                    elif level == 3:
                        self.sub_section = section.group().strip("# ")
                        # self.sub_section = re.sub(r"[^\d\.\-]", "", sub_section)
                content += line + "\n"
            else:
                if self.chapter:
                    content += line + "\n"

        self.save_result(save_dir)

    def save_result(self, save_dir="text_blocks"):
        filename = os.path.splitext(os.path.basename(self.book_path))[0] + ".json"
        save_path = os.path.join(save_dir, filename)
        os.makedirs("text_blocks", exist_ok=True)
        with open(save_path, "w", encoding="utf-8") as f:
            json.dump(self.output, f, ensure_ascii=False, indent=4)


class CodeSplitter(BaseSplitter):
    def __init__(self, code_path) -> None:
        super().__init__(code_path)
        self.LEN_LIMIT_HIGH = 1500
        self.LEN_LIMIT_LOW = 500

    def define_regex(self):
        self.chapter_pattern = re.compile(r"^## \d+ .+?$", re.M)
        self.section_pattern = re.compile(r"^## [\d\. ]{3,}.+?$", re.M)
        self.appendix_pattern = re.compile(r"^## 附录 .+?$", re.M)
        self.title_pattern = re.compile(r"^## .+?$", re.M)

    def process_book(self, save_dir):
        content = ""
        self.chapter, self.section, self.sub_section = "", "", ""
        text = self._preprocess(self.text)
        # text = self.text
        text_list = text.split("\n")
        for i, line in enumerate(text_list):
            if line.strip() == "":
                continue
            if chapter := self.chapter_pattern.search(line):
                if self.chapter:
                    # a new chapter starts, save the previous content
                    self.add_new_content(content)
                    content = ""

                self.chapter = chapter.group().strip("# ")
                self.section, self.sub_section = "", ""
            elif section := self.section_pattern.search(line):
                if len(content) > self.LEN_LIMIT_LOW:
                    self.add_new_content(content)
                    content = ""
                content += line + "\n"

                level = self.determine_level(line)
                if level == 2:
                    self.section = section.group().strip("# ")
                    # self.section = re.sub(r"[^\d\.\-]", "", section)
                    self.sub_section = ""
                elif level == 3:
                    self.sub_section = section.group().strip("# ")

            elif appendix := self.appendix_pattern.search(line):
                if len(content) > self.LEN_LIMIT_LOW:
                    self.add_new_content(content)
                    content = line + "\n"
                self.chapter = appendix.group().strip("# ")
                self.section, self.sub_section = "", ""

            elif self.title_pattern.search(line):
                if self.chapter:
                    if len(content) > self.LEN_LIMIT_LOW:
                        self.add_new_content(content)
                        content = line + "\n"

            else:
                if self.chapter:
                    content += line + "\n"

        self.add_new_content(content)

        self.save_result(save_dir)


if __name__ == "__main__":
    save_dir = r"text_blocks\codes"

    for fname in os.listdir(r"raw_data\codes"):
        if fname.endswith(".md"):
            book_path = os.path.join(r"raw_data\codes", fname)
            splitter = CodeSplitter(book_path)
            splitter.process_book(save_dir)
    print("Done!")
